"""
因子分析模块

提供因子挖掘、构建、验证和组合优化功能，支持技术因子、基本面因子和另类因子。
"""

from typing import Dict, List, Any, Optional, Union, Tuple
from datetime import datetime, timedelta
import asyncio
import numpy as np
import pandas as pd
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum

from ..config.settings import settings


class FactorType(Enum):
    """因子类型枚举"""
    TECHNICAL = "technical"          # 技术因子
    FUNDAMENTAL = "fundamental"      # 基本面因子
    ALTERNATIVE = "alternative"      # 另类因子
    MACRO = "macro"                 # 宏观因子
    SENTIMENT = "sentiment"         # 情绪因子


class FactorCategory(Enum):
    """因子分类枚举"""
    MOMENTUM = "momentum"           # 动量因子
    REVERSAL = "reversal"          # 反转因子
    VALUE = "value"                # 价值因子
    GROWTH = "growth"              # 成长因子
    QUALITY = "quality"            # 质量因子
    VOLATILITY = "volatility"      # 波动率因子
    LIQUIDITY = "liquidity"        # 流动性因子
    SIZE = "size"                  # 规模因子


@dataclass
class FactorMetrics:
    """因子评价指标"""
    ic: float                      # 信息系数
    ic_ir: float                   # IC信息比率
    rank_ic: float                 # 排序IC
    rank_ic_ir: float             # 排序IC信息比率
    turnover: float               # 换手率
    max_drawdown: float           # 最大回撤
    sharpe_ratio: float           # 夏普比率
    calmar_ratio: float           # 卡玛比率
    win_rate: float               # 胜率
    factor_coverage: float        # 因子覆盖度
    factor_decay: float           # 因子衰减度
    
    def to_dict(self) -> Dict[str, float]:
        """转换为字典"""
        return {
            "ic": self.ic,
            "ic_ir": self.ic_ir,
            "rank_ic": self.rank_ic,
            "rank_ic_ir": self.rank_ic_ir,
            "turnover": self.turnover,
            "max_drawdown": self.max_drawdown,
            "sharpe_ratio": self.sharpe_ratio,
            "calmar_ratio": self.calmar_ratio,
            "win_rate": self.win_rate,
            "factor_coverage": self.factor_coverage,
            "factor_decay": self.factor_decay
        }


@dataclass
class FactorInfo:
    """因子信息"""
    factor_id: str
    name: str
    description: str
    factor_type: FactorType
    category: FactorCategory
    formula: str
    parameters: Dict[str, Any]
    created_at: datetime
    updated_at: datetime
    metrics: Optional[FactorMetrics] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "factor_id": self.factor_id,
            "name": self.name,
            "description": self.description,
            "factor_type": self.factor_type.value,
            "category": self.category.value,
            "formula": self.formula,
            "parameters": self.parameters,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "metrics": self.metrics.to_dict() if self.metrics else None
        }


class BaseFactor(ABC):
    """基础因子类"""
    
    def __init__(
        self,
        factor_id: str,
        name: str,
        description: str,
        factor_type: FactorType,
        category: FactorCategory,
        parameters: Optional[Dict[str, Any]] = None
    ):
        self.factor_id = factor_id
        self.name = name
        self.description = description
        self.factor_type = factor_type
        self.category = category
        self.parameters = parameters or {}
        self.created_at = datetime.now()
        self.updated_at = datetime.now()
    
    @abstractmethod
    async def calculate(
        self,
        data: Dict[str, pd.DataFrame],
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """
        计算因子值
        
        Args:
            data: 输入数据字典
            start_date: 开始日期
            end_date: 结束日期
            
        Returns:
            因子值DataFrame
        """
        pass
    
    @abstractmethod
    def get_formula(self) -> str:
        """获取因子公式"""
        pass
    
    def get_info(self) -> FactorInfo:
        """获取因子信息"""
        return FactorInfo(
            factor_id=self.factor_id,
            name=self.name,
            description=self.description,
            factor_type=self.factor_type,
            category=self.category,
            formula=self.get_formula(),
            parameters=self.parameters,
            created_at=self.created_at,
            updated_at=self.updated_at
        )


class TechnicalFactor(BaseFactor):
    """技术因子基类"""
    
    def __init__(self, factor_id: str, name: str, description: str, category: FactorCategory, **kwargs):
        super().__init__(factor_id, name, description, FactorType.TECHNICAL, category, kwargs)


class MomentumFactor(TechnicalFactor):
    """动量因子"""
    
    def __init__(self, lookback_period: int = 20):
        super().__init__(
            factor_id="momentum",
            name="动量因子",
            description="基于价格动量的技术因子",
            category=FactorCategory.MOMENTUM,
            lookback_period=lookback_period
        )
    
    async def calculate(
        self,
        data: Dict[str, pd.DataFrame],
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """计算动量因子"""
        price_data = data.get("price", pd.DataFrame())
        if price_data.empty:
            return pd.DataFrame()
        
        lookback = self.parameters.get("lookback_period", 20)
        
        # 计算收益率
        returns = price_data.pct_change()
        
        # 计算动量（过去N期累计收益率）
        momentum = returns.rolling(window=lookback).sum()
        
        return momentum.loc[start_date:end_date]
    
    def get_formula(self) -> str:
        return f"sum(returns, {self.parameters.get('lookback_period', 20)})"


class ReversalFactor(TechnicalFactor):
    """反转因子"""
    
    def __init__(self, lookback_period: int = 5):
        super().__init__(
            factor_id="reversal",
            name="反转因子",
            description="基于短期反转的技术因子",
            category=FactorCategory.REVERSAL,
            lookback_period=lookback_period
        )
    
    async def calculate(
        self,
        data: Dict[str, pd.DataFrame],
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """计算反转因子"""
        price_data = data.get("price", pd.DataFrame())
        if price_data.empty:
            return pd.DataFrame()
        
        lookback = self.parameters.get("lookback_period", 5)
        
        # 计算收益率
        returns = price_data.pct_change()
        
        # 计算反转因子（短期收益率的负值）
        reversal = -returns.rolling(window=lookback).sum()
        
        return reversal.loc[start_date:end_date]
    
    def get_formula(self) -> str:
        return f"-sum(returns, {self.parameters.get('lookback_period', 5)})"


class VolatilityFactor(TechnicalFactor):
    """波动率因子"""
    
    def __init__(self, lookback_period: int = 20):
        super().__init__(
            factor_id="volatility",
            name="波动率因子",
            description="基于价格波动率的技术因子",
            category=FactorCategory.VOLATILITY,
            lookback_period=lookback_period
        )
    
    async def calculate(
        self,
        data: Dict[str, pd.DataFrame],
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """计算波动率因子"""
        price_data = data.get("price", pd.DataFrame())
        if price_data.empty:
            return pd.DataFrame()
        
        lookback = self.parameters.get("lookback_period", 20)
        
        # 计算收益率
        returns = price_data.pct_change()
        
        # 计算滚动标准差
        volatility = returns.rolling(window=lookback).std()
        
        return volatility.loc[start_date:end_date]
    
    def get_formula(self) -> str:
        return f"std(returns, {self.parameters.get('lookback_period', 20)})"


class FundamentalFactor(BaseFactor):
    """基本面因子基类"""
    
    def __init__(self, factor_id: str, name: str, description: str, category: FactorCategory, **kwargs):
        super().__init__(factor_id, name, description, FactorType.FUNDAMENTAL, category, kwargs)


class ValueFactor(FundamentalFactor):
    """价值因子"""
    
    def __init__(self):
        super().__init__(
            factor_id="pe_ratio",
            name="市盈率因子",
            description="基于市盈率的价值因子",
            category=FactorCategory.VALUE
        )
    
    async def calculate(
        self,
        data: Dict[str, pd.DataFrame],
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """计算价值因子"""
        fundamental_data = data.get("fundamental", pd.DataFrame())
        if fundamental_data.empty:
            return pd.DataFrame()
        
        # 计算市盈率倒数作为价值因子
        pe_ratio = fundamental_data.get("pe_ratio", pd.Series())
        value_factor = 1 / pe_ratio.replace(0, np.nan)
        
        return value_factor.to_frame().loc[start_date:end_date]
    
    def get_formula(self) -> str:
        return "1 / pe_ratio"


class GrowthFactor(FundamentalFactor):
    """成长因子"""
    
    def __init__(self):
        super().__init__(
            factor_id="revenue_growth",
            name="营收增长因子",
            description="基于营收增长率的成长因子",
            category=FactorCategory.GROWTH
        )
    
    async def calculate(
        self,
        data: Dict[str, pd.DataFrame],
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """计算成长因子"""
        fundamental_data = data.get("fundamental", pd.DataFrame())
        if fundamental_data.empty:
            return pd.DataFrame()
        
        # 计算营收增长率
        revenue = fundamental_data.get("revenue", pd.Series())
        growth_factor = revenue.pct_change(periods=4)  # 年度增长率
        
        return growth_factor.to_frame().loc[start_date:end_date]
    
    def get_formula(self) -> str:
        return "revenue.pct_change(4)"


class FactorAnalyzer:
    """因子分析器"""
    
    def __init__(self):
        """初始化因子分析器"""
        self._factors: Dict[str, BaseFactor] = {}
        self._factor_data: Dict[str, pd.DataFrame] = {}
        self._return_data: Optional[pd.DataFrame] = None
        self._initialized = False
    
    async def initialize(self) -> None:
        """初始化因子分析器"""
        if self._initialized:
            return
        
        # 注册默认因子
        await self._register_default_factors()
        
        self._initialized = True
    
    async def _register_default_factors(self) -> None:
        """注册默认因子"""
        # 技术因子
        momentum_factor = MomentumFactor(lookback_period=20)
        reversal_factor = ReversalFactor(lookback_period=5)
        volatility_factor = VolatilityFactor(lookback_period=20)
        
        # 基本面因子
        value_factor = ValueFactor()
        growth_factor = GrowthFactor()
        
        # 注册因子
        factors = [momentum_factor, reversal_factor, volatility_factor, value_factor, growth_factor]
        for factor in factors:
            self._factors[factor.factor_id] = factor
    
    async def register_factor(self, factor: BaseFactor) -> None:
        """
        注册因子
        
        Args:
            factor: 因子实例
        """
        self._factors[factor.factor_id] = factor
        factor.updated_at = datetime.now()
    
    async def calculate_factor(
        self,
        factor_id: str,
        data: Dict[str, pd.DataFrame],
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """
        计算单个因子
        
        Args:
            factor_id: 因子ID
            data: 输入数据
            start_date: 开始日期
            end_date: 结束日期
            
        Returns:
            因子值DataFrame
        """
        if not self._initialized:
            await self.initialize()
        
        factor = self._factors.get(factor_id)
        if not factor:
            raise ValueError(f"Factor {factor_id} not found")
        
        factor_values = await factor.calculate(data, start_date, end_date)
        
        # 缓存因子数据
        self._factor_data[factor_id] = factor_values
        
        return factor_values
    
    async def calculate_all_factors(
        self,
        data: Dict[str, pd.DataFrame],
        start_date: datetime,
        end_date: datetime,
        factor_ids: Optional[List[str]] = None
    ) -> Dict[str, pd.DataFrame]:
        """
        计算所有因子
        
        Args:
            data: 输入数据
            start_date: 开始日期
            end_date: 结束日期
            factor_ids: 指定因子ID列表
            
        Returns:
            因子值字典
        """
        if not self._initialized:
            await self.initialize()
        
        target_factors = factor_ids or list(self._factors.keys())
        results = {}
        
        # 并行计算因子
        tasks = []
        for factor_id in target_factors:
            if factor_id in self._factors:
                task = self.calculate_factor(factor_id, data, start_date, end_date)
                tasks.append((factor_id, task))
        
        # 等待所有任务完成
        for factor_id, task in tasks:
            try:
                factor_values = await task
                results[factor_id] = factor_values
            except Exception as e:
                print(f"Error calculating factor {factor_id}: {e}")
        
        return results
    
    async def evaluate_factor(
        self,
        factor_id: str,
        factor_data: pd.DataFrame,
        return_data: pd.DataFrame,
        periods: List[int] = [1, 5, 10, 20]
    ) -> FactorMetrics:
        """
        评价因子有效性
        
        Args:
            factor_id: 因子ID
            factor_data: 因子数据
            return_data: 收益率数据
            periods: 预测期间列表
            
        Returns:
            因子评价指标
        """
        if factor_data.empty or return_data.empty:
            return FactorMetrics(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
        
        # 对齐数据
        aligned_factor, aligned_return = factor_data.align(return_data, join='inner')
        
        if aligned_factor.empty or aligned_return.empty:
            return FactorMetrics(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
        
        # 计算各期间的IC
        ic_results = []
        rank_ic_results = []
        
        for period in periods:
            # 计算未来收益率
            future_returns = aligned_return.shift(-period)
            
            # 计算IC（信息系数）
            ic_series = []
            rank_ic_series = []
            
            for date in aligned_factor.index:
                if date in future_returns.index:
                    factor_values = aligned_factor.loc[date].dropna()
                    return_values = future_returns.loc[date].dropna()
                    
                    # 对齐数据
                    common_symbols = factor_values.index.intersection(return_values.index)
                    if len(common_symbols) > 10:  # 至少需要10个样本
                        factor_vals = factor_values.loc[common_symbols]
                        return_vals = return_values.loc[common_symbols]
                        
                        # 计算相关系数
                        ic = factor_vals.corr(return_vals)
                        rank_ic = factor_vals.rank().corr(return_vals.rank())
                        
                        if not np.isnan(ic):
                            ic_series.append(ic)
                        if not np.isnan(rank_ic):
                            rank_ic_series.append(rank_ic)
            
            if ic_series:
                ic_results.extend(ic_series)
            if rank_ic_series:
                rank_ic_results.extend(rank_ic_series)
        
        # 计算统计指标
        ic_mean = np.mean(ic_results) if ic_results else 0
        ic_std = np.std(ic_results) if ic_results else 0
        ic_ir = ic_mean / ic_std if ic_std > 0 else 0
        
        rank_ic_mean = np.mean(rank_ic_results) if rank_ic_results else 0
        rank_ic_std = np.std(rank_ic_results) if rank_ic_results else 0
        rank_ic_ir = rank_ic_mean / rank_ic_std if rank_ic_std > 0 else 0
        
        # 计算换手率
        turnover = self._calculate_turnover(aligned_factor)
        
        # 计算其他指标
        max_drawdown = self._calculate_max_drawdown(aligned_factor)
        sharpe_ratio = self._calculate_sharpe_ratio(aligned_factor, aligned_return)
        calmar_ratio = sharpe_ratio / max_drawdown if max_drawdown > 0 else 0
        win_rate = len([ic for ic in ic_results if ic > 0]) / len(ic_results) if ic_results else 0
        factor_coverage = aligned_factor.notna().mean().mean()
        factor_decay = self._calculate_factor_decay(aligned_factor)
        
        return FactorMetrics(
            ic=ic_mean,
            ic_ir=ic_ir,
            rank_ic=rank_ic_mean,
            rank_ic_ir=rank_ic_ir,
            turnover=turnover,
            max_drawdown=max_drawdown,
            sharpe_ratio=sharpe_ratio,
            calmar_ratio=calmar_ratio,
            win_rate=win_rate,
            factor_coverage=factor_coverage,
            factor_decay=factor_decay
        )
    
    def _calculate_turnover(self, factor_data: pd.DataFrame) -> float:
        """计算因子换手率"""
        if factor_data.empty:
            return 0
        
        # 计算因子排名
        factor_ranks = factor_data.rank(axis=1, pct=True)
        
        # 计算相邻期间排名变化
        rank_changes = factor_ranks.diff().abs()
        
        # 计算平均换手率
        turnover = rank_changes.mean().mean()
        
        return turnover if not np.isnan(turnover) else 0
    
    def _calculate_max_drawdown(self, factor_data: pd.DataFrame) -> float:
        """计算最大回撤"""
        if factor_data.empty:
            return 0
        
        # 计算因子累计值
        factor_cumsum = factor_data.mean(axis=1).cumsum()
        
        # 计算回撤
        peak = factor_cumsum.expanding().max()
        drawdown = (factor_cumsum - peak) / peak
        
        max_dd = drawdown.min()
        
        return abs(max_dd) if not np.isnan(max_dd) else 0
    
    def _calculate_sharpe_ratio(self, factor_data: pd.DataFrame, return_data: pd.DataFrame) -> float:
        """计算夏普比率"""
        if factor_data.empty or return_data.empty:
            return 0
        
        # 构建因子组合收益率
        factor_weights = factor_data.rank(axis=1, pct=True)
        portfolio_returns = (factor_weights * return_data).sum(axis=1)
        
        # 计算夏普比率
        mean_return = portfolio_returns.mean()
        std_return = portfolio_returns.std()
        
        sharpe = mean_return / std_return if std_return > 0 else 0
        
        return sharpe if not np.isnan(sharpe) else 0
    
    def _calculate_factor_decay(self, factor_data: pd.DataFrame) -> float:
        """计算因子衰减度"""
        if factor_data.empty:
            return 0
        
        # 计算因子自相关性
        factor_mean = factor_data.mean(axis=1)
        autocorr = factor_mean.autocorr(lag=1)
        
        # 衰减度 = 1 - 自相关性
        decay = 1 - autocorr if not np.isnan(autocorr) else 0
        
        return max(0, decay)
    
    async def optimize_factor_combination(
        self,
        factor_data: Dict[str, pd.DataFrame],
        return_data: pd.DataFrame,
        method: str = "ic_weighted",
        constraints: Optional[Dict[str, Any]] = None
    ) -> Dict[str, float]:
        """
        优化因子组合
        
        Args:
            factor_data: 因子数据字典
            return_data: 收益率数据
            method: 优化方法
            constraints: 约束条件
            
        Returns:
            因子权重字典
        """
        if not factor_data or return_data.empty:
            return {}
        
        # 计算各因子的IC
        factor_ics = {}
        for factor_id, data in factor_data.items():
            metrics = await self.evaluate_factor(factor_id, data, return_data)
            factor_ics[factor_id] = abs(metrics.ic)
        
        if method == "ic_weighted":
            # 基于IC加权
            total_ic = sum(factor_ics.values())
            if total_ic > 0:
                weights = {fid: ic / total_ic for fid, ic in factor_ics.items()}
            else:
                # 等权重
                n_factors = len(factor_ics)
                weights = {fid: 1.0 / n_factors for fid in factor_ics.keys()}
        
        elif method == "equal_weight":
            # 等权重
            n_factors = len(factor_ics)
            weights = {fid: 1.0 / n_factors for fid in factor_ics.keys()}
        
        else:
            # 默认等权重
            n_factors = len(factor_ics)
            weights = {fid: 1.0 / n_factors for fid in factor_ics.keys()}
        
        return weights
    
    async def get_factor_info(self, factor_id: str) -> Optional[FactorInfo]:
        """
        获取因子信息
        
        Args:
            factor_id: 因子ID
            
        Returns:
            因子信息
        """
        if not self._initialized:
            await self.initialize()
        
        factor = self._factors.get(factor_id)
        if not factor:
            return None
        
        info = factor.get_info()
        
        # 如果有缓存的因子数据，计算评价指标
        if factor_id in self._factor_data and self._return_data is not None:
            metrics = await self.evaluate_factor(
                factor_id,
                self._factor_data[factor_id],
                self._return_data
            )
            info.metrics = metrics
        
        return info
    
    async def list_factors(
        self,
        factor_type: Optional[FactorType] = None,
        category: Optional[FactorCategory] = None
    ) -> List[FactorInfo]:
        """
        列出因子
        
        Args:
            factor_type: 因子类型筛选
            category: 因子分类筛选
            
        Returns:
            因子信息列表
        """
        if not self._initialized:
            await self.initialize()
        
        factor_infos = []
        
        for factor_id, factor in self._factors.items():
            # 应用筛选条件
            if factor_type and factor.factor_type != factor_type:
                continue
            if category and factor.category != category:
                continue
            
            info = await self.get_factor_info(factor_id)
            if info:
                factor_infos.append(info)
        
        return factor_infos
    
    async def remove_factor(self, factor_id: str) -> bool:
        """
        移除因子
        
        Args:
            factor_id: 因子ID
            
        Returns:
            是否成功移除
        """
        if factor_id in self._factors:
            del self._factors[factor_id]
            
            # 清理缓存数据
            if factor_id in self._factor_data:
                del self._factor_data[factor_id]
            
            return True
        
        return False
    
    async def set_return_data(self, return_data: pd.DataFrame) -> None:
        """
        设置收益率数据
        
        Args:
            return_data: 收益率数据
        """
        self._return_data = return_data
    
    async def get_factor_correlation_matrix(
        self,
        factor_ids: Optional[List[str]] = None
    ) -> pd.DataFrame:
        """
        获取因子相关性矩阵
        
        Args:
            factor_ids: 因子ID列表
            
        Returns:
            相关性矩阵
        """
        if not self._factor_data:
            return pd.DataFrame()
        
        target_factors = factor_ids or list(self._factor_data.keys())
        
        # 构建因子数据矩阵
        factor_matrix = pd.DataFrame()
        for factor_id in target_factors:
            if factor_id in self._factor_data:
                factor_data = self._factor_data[factor_id]
                if not factor_data.empty:
                    # 使用因子的平均值
                    factor_series = factor_data.mean(axis=1)
                    factor_matrix[factor_id] = factor_series
        
        if factor_matrix.empty:
            return pd.DataFrame()
        
        # 计算相关性矩阵
        correlation_matrix = factor_matrix.corr()
        
        return correlation_matrix
    
    async def cleanup(self) -> None:
        """清理资源"""
        self._factor_data.clear()
        self._return_data = None
        self._initialized = False